// Copyright (c) 2020 Presto Labs Pte. Ltd.
// Author: taekwon

#pragma once

#include <atomic>
#include <functional>
#include <list>
#include <map>
#include <memory>
#include <string>
#include <utility>

#include "coin/proto/coin_telemetry.pb.h"
#include "coin2/base/time.h"
#include "coin2/exchange/base/order/order.h"
#include "presto/quant/math/sampled_moving_average.h"

namespace coin2::exchange::base::strategy_util {

struct LatencyWindow {
  typedef std::pair<MarketExchangeApi, coin::proto::StrategyComponentLatencyStatsProto::LatencyType>
      LatencyGroup;
  std::map<LatencyGroup, double> latency_window;
  std::map<LatencyGroup, int64_t> last_flush_ts;

  bool UpdateLatencyWindow(
      int64_t timestamp,
      const MarketExchangeApi& mea,
      coin::proto::StrategyComponentLatencyStatsProto::LatencyType lat_type,
      int64_t latency,
      double* twindow_latency_sum,
      double* twindow_length) {
    auto key = LatencyGroup{mea, lat_type};
    if (latency_window.count(key) == 0) {
      latency_window[key] = 0;
      last_flush_ts[key] = timestamp;
    }
    latency_window[key] += latency;
    if (last_flush_ts[key] + 600e9 < timestamp) {
      *twindow_latency_sum = latency_window[key];
      *twindow_length = timestamp - last_flush_ts[key];
      latency_window[key] = 0;
      last_flush_ts[key] = timestamp;
      return true;
    } else {
      return false;
    }
  }
};  // struct LatencyWindow

class LatencyRecorder {
 public:
  using OrderContext = coin2::exchange::base::order::IOrderContext;

 public:
  typedef int64_t key_type;
  typedef coin::proto::LatencyProto value_type;
  typedef std::list<key_type> list_type;
  typedef std::map<key_type, std::pair<std::unique_ptr<value_type>, typename list_type::iterator>>
      map_type;

 public:
  LatencyRecorder(
      size_t cap,
      std::function<bool(const coin::proto::LatencyProto&, bool, int64_t)> write_func)
      : cap_(cap), write_func_(std::move(write_func)), context_id_(GetCurrentTimestamp() / 1000) {}

  void Begin(coin::proto::LatencyProto::LatencyTag tag);

  void End(int64_t context_id);

  void SetTimeoutMark(const coin::proto::AccountRequestProto& acc_request, const std::string& tag);

  void SetBeforeSubmitMark(
      const OrderContext* const oc,
      const coin::proto::AccountRequestProto& acc_request);

  void SetAfterSubmitMark(
      const OrderContext* const oc,
      const coin::proto::AccountRequestProto& acc_request);

  void SetBeforeCancelMark(
      const OrderContext* const oc,
      const coin::proto::AccountRequestProto& acc_request);

  void SetAfterCancelMark(
      const OrderContext* const oc,
      const coin::proto::AccountRequestProto& acc_request);

  void SetBeforeAmendMark(
      const OrderContext* const oc,
      const coin::proto::AccountRequestProto& acc_request);

  void SetAfterAmendMark(
      const OrderContext* const oc,
      const coin::proto::AccountRequestProto& acc_request);

  void SetProtoAndTime(
      int64_t context_id,
      const coin::proto::LifeOfSignal& point_name,
      int64_t timestamp,
      const coin::proto::AccountRequestProto& account_request,
      const std::string& tag = "");

  int64_t GetCurrentContextId() const { return context_id_; }

  void WriteTps(const coin::proto::TpsProto& tps);

  void set_write_telemetry(
      std::function<bool(const coin::proto::TelemetryProto&, bool, int64_t)>&& write_telemetry) {
    write_telemetry_ = std::move(write_telemetry);
  }

  LatencyWindow* GetMutableLatencyWindow() { return &latency_window_; }

 private:
  int64_t GetNextContextId() { return ++context_id_; }

  int64_t GetTimestamp() {
    int64_t curr_ts = (curr_ts_.has_value()) ? curr_ts_.value() : GetCurrentTimestamp();
    return curr_ts;
  }

  void SetTimestamp(int64_t curr_ts) { curr_ts_ = curr_ts; }

  // From <boost/compute/detail/lru_cache.hpp>
  size_t size() const { return m_map.size(); }

  void PutIntoCache(const key_type& key, std::unique_ptr<value_type> value);

  map_type::iterator GetFromCache(const key_type& key);

  void WriteTelemetry(map_type::const_iterator i);

  bool IsLatencyEvicted(const value_type* latency, bool ignore_pending_order);

  bool Evict(bool ignore_pending_order = false);

  const map_type* GetMmap() const { return &m_map; }

  const list_type* GetMlist() const { return &m_list; }

 private:
  const size_t cap_;
  std::function<bool(const coin::proto::LatencyProto&, bool, int64_t)> write_func_;
  std::atomic<int64_t> context_id_;
  std::map<coin::proto::LatencyProto::LatencyTag, int64_t> prev_feed_latency_send_ts_;
  std::optional<int64_t> curr_ts_;

  map_type m_map;
  list_type m_list;

  // For tps proto
  coin::proto::TelemetryProto telemetry_;
  std::function<bool(const coin::proto::TelemetryProto&, bool, int64_t)> write_telemetry_;
  LatencyWindow latency_window_;
  friend class LatencyRecorderTest;
  friend class LatencyRecorderGlobal;
};

}  // namespace coin2::exchange::base::strategy_util
